import numpy as np
from numpy import array, mean, where
from numpy.linalg import norm
import cv2
import math
from tqdm import tqdm
import matplotlib.pyplot as plt

IMG_PATH = "lenna.png"

img = cv2.imread(IMG_PATH)
img = cv2.GaussianBlur(img, ksize=(3,3), sigmaX=0)
lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)

# 总像素数
N = img.shape[0]*img.shape[1]
# 聚类数
K = 1000
# 每个超像素的初始大小
S = int(math.sqrt(N/K))
# 距离权重
M = 20
# 迭代次数
ITER = 5
# 距离函数
def distance(pix1_lab, pix1_loc, pix2_lab, pix2_loc, m, s):
    dlab = norm(array(pix1_lab, dtype='int64')-array(pix2_lab, dtype='int64'))
    dxy = norm(pix1_loc-pix2_loc)
    dist = dlab**2 + (m/s*dxy)**2
    #print(dlab, dxy, dist)
    return dist
    
# 设置初始采样点
sample = []
for x in range(int(S/2),img.shape[0],int(S)):
    for y in range(int(S/2),img.shape[1],int(S)):
        sample.append([x,y])

# 计算原图像的梯度
grad = np.zeros((img.shape[0], img.shape[1]))
for x in range(1, grad.shape[0]-1):
    for y in range(1, grad.shape[1]-1):
        grad_x = np.linalg.norm(lab[x+1,y]-lab[x-1,y])
        grad_y = np.linalg.norm(lab[x,+1]-lab[x,y-1])
        grad[x,y] = grad_x+grad_y

# 移动采样点，远离梯度
adjusted_sample = []
for each in sample:
    search_size = 2
    search_zone = grad[each[0]-search_size:each[0]+search_size, each[1]-search_size:each[1]+search_size]
    lowest_grad = np.amin(search_zone)
    t = np.where(search_zone==lowest_grad)
    new_sample = [each[0]-search_size+t[0][0], each[1]-search_size+t[1][0]]
    adjusted_sample.append(new_sample)
    
# sample存储的是位置
np_sample = np.array(adjusted_sample)
sup_pixel_sample = np.hstack((np_sample, np.arange(len(np_sample)).reshape((len(np_sample),1))))

import time
for itration in range(ITER):
    # 开始聚类
    sup_pixel_frame = np.zeros((img.shape[0], img.shape[1]))

    frame = lab.copy()
    for x in tqdm(range(img.shape[0])):
        for y in range(img.shape[1]):
            sample_index = where((sup_pixel_sample[:,0]<x+2*S)&(sup_pixel_sample[:,0]>x-2*S)&(sup_pixel_sample[:,1]<y+2*S)&(sup_pixel_sample[:,1]>y-2*S))[0]
            indicate = sup_pixel_sample[sample_index]
            point_loc = array([x,y])
            point_lab = lab[x][y]
            min_dist = 999999
            for indicate_loc in indicate:
                indicate_lab = lab[indicate_loc[0]][indicate_loc[1]]
                dist = distance(point_lab, point_loc, indicate_lab, indicate_loc[:2], M, S)
                #time.sleep(0.5)
                if dist<min_dist:
                    result_loc = indicate_loc
                    min_dist = dist
                #print(point_lab, point_loc, indicate_lab, indicate_loc[:2], dist)
            frame[x][y] = lab[result_loc[0]][result_loc[1]]
            sup_pixel_frame[x][y] = result_loc[2]
            #print(result_loc,'---------------')


            
    #重新设定中心
    iterate_sample = []
    for i in range(len(sup_pixel_sample)):
        x_index, y_index = where(sup_pixel_frame == i)
        try:
            x = int(mean(x_index))
            y = int(mean(y_index))
        except:
            x = -1000
            y = -1000
        iterate_sample.append([x, y, sup_pixel_sample[i][2]])
    sup_pixel_sample = np.array(iterate_sample)

frame2 = cv2.cvtColor(frame, cv2.COLOR_LAB2BGR)
cv2.imwrite('result_K'+str(K)+'_M'+str(M)+'.jpg', frame2)